#!/usr/bin/env python3
"""
Claude Code <-> Copilot 映射代理服务器
实现透明的请求转发和响应转换
"""

import asyncio
import json
import logging
import time
from typing import Dict, Any, Optional
from aiohttp import web, ClientSession, ClientTimeout
from aiohttp.web_request import Request
from aiohttp.web_response import Response, StreamResponse
import aiohttp_cors

from request_transformer import RequestTransformer


class MappingProxy:
    def __init__(self, 
                 listen_host: str = "127.0.0.1",
                 listen_port: int = 8080,
                 copilot_base_url: str = "https://api.business.githubcopilot.com"):
        """初始化代理服务器"""
        self.listen_host = listen_host
        self.listen_port = listen_port  
        self.copilot_base_url = copilot_base_url
        self.transformer = RequestTransformer()
        self.session: Optional[ClientSession] = None
        
        # 配置日志
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(levelname)s - %(message)s'
        )
        self.logger = logging.getLogger(__name__)
    
    async def create_app(self) -> web.Application:
        """创建web应用"""
        app = web.Application()
        
        # 配置CORS
        cors = aiohttp_cors.setup(app, defaults={
            "*": aiohttp_cors.ResourceOptions(
                allow_credentials=True,
                expose_headers="*",
                allow_headers="*",
                allow_methods="*"
            )
        })
        
        # 注册路由
        app.router.add_route("POST", "/v1/messages", self.handle_claude_request)
        app.router.add_route("GET", "/health", self.health_check)
        app.router.add_route("GET", "/status", self.get_status)
        
        # 添加CORS到所有路由
        for route in list(app.router.routes()):
            cors.add(route)
        
        return app
    
    async def start_server(self):
        """启动代理服务器"""
        self.session = ClientSession(
            timeout=ClientTimeout(total=300, connect=30),
            connector=aiohttp.TCPConnector(limit=100, limit_per_host=30)
        )
        
        app = await self.create_app()
        runner = web.AppRunner(app)
        await runner.setup()
        
        site = web.TCPSite(runner, self.listen_host, self.listen_port)
        await site.start()
        
        self.logger.info(f"Claude Code -> Copilot 代理服务器启动于 http://{self.listen_host}:{self.listen_port}")
        self.logger.info(f"转发目标: {self.copilot_base_url}")
        
        return runner
    
    async def handle_claude_request(self, request: Request) -> Response:
        """处理Claude Code请求"""
        start_time = time.time()
        client_ip = request.remote
        
        try:
            # 解析原始请求
            claude_request = await self._parse_claude_request(request)
            self.logger.info(f"收到Claude请求来自 {client_ip}: {claude_request['method']} {claude_request['url']}")
            
            # 转换为Copilot格式
            copilot_request = self.transformer.claude_to_copilot(claude_request)
            self.logger.debug(f"转换为Copilot请求: {json.dumps(copilot_request, indent=2)}")
            
            # 发送到Copilot API
            copilot_response = await self._forward_to_copilot(copilot_request)
            
            # 转换响应为Claude格式
            claude_response = await self._transform_copilot_response(copilot_response)
            
            # 记录处理时间
            process_time = time.time() - start_time
            self.logger.info(f"请求处理完成，耗时: {process_time:.2f}秒")
            
            return claude_response
            
        except Exception as e:
            self.logger.error(f"处理请求出错: {str(e)}", exc_info=True)
            return await self._create_error_response(str(e))
    
    async def _parse_claude_request(self, request: Request) -> Dict[str, Any]:
        """解析Claude请求"""
        body_text = await request.text()
        
        claude_request = {
            'method': request.method,
            'url': str(request.url.with_host('api.anthropic.com').with_scheme('https')),
            'headers': dict(request.headers),
            'body': body_text
        }
        
        return claude_request
    
    async def _forward_to_copilot(self, copilot_request: Dict[str, Any]) -> Dict[str, Any]:
        """转发请求到Copilot API"""
        url = copilot_request['url']
        headers = copilot_request['headers']
        body = copilot_request['body']
        
        # 确保请求体是JSON字符串
        if isinstance(body, dict):
            body_json = json.dumps(body)
        else:
            body_json = body
            
        headers['content-length'] = str(len(body_json.encode('utf-8')))
        
        self.logger.debug(f"发送到Copilot: {url}")
        self.logger.debug(f"请求头: {json.dumps(headers, indent=2)}")
        
        async with self.session.post(url, headers=headers, data=body_json) as resp:
            response_headers = dict(resp.headers)
            response_status = resp.status
            
            # 处理流式响应
            if response_headers.get('content-type', '').startswith('text/event-stream'):
                response_body = await self._handle_sse_response(resp)
            else:
                response_body = await resp.text()
            
            return {
                'status': response_status,
                'headers': response_headers,
                'body': response_body
            }
    
    async def _handle_sse_response(self, resp) -> Dict[str, Any]:
        """处理服务器发送事件(SSE)响应"""
        events = []
        current_event = {}
        
        async for line in resp.content:
            line = line.decode('utf-8').strip()
            
            if line.startswith('event:'):
                if current_event:
                    events.append(current_event.copy())
                current_event = {'type': line[6:].strip()}
            elif line.startswith('data:'):
                data_str = line[5:].strip()
                if data_str == '[DONE]':
                    break
                try:
                    current_event['data'] = json.loads(data_str)
                except:
                    current_event['data'] = data_str
            elif line == '' and current_event:
                events.append(current_event.copy())
                current_event = {}
        
        if current_event:
            events.append(current_event)
        
        return {
            'format': 'sse',
            'events': events
        }
    
    async def _transform_copilot_response(self, copilot_response: Dict[str, Any]) -> Response:
        """转换Copilot响应为Claude格式"""
        status = copilot_response.get('status', 200)
        body = copilot_response.get('body', {})
        
        # 处理SSE到JSON的转换
        if isinstance(body, dict) and body.get('format') == 'sse':
            claude_json = self._convert_sse_to_claude_json(body['events'])
        else:
            # 处理普通JSON响应
            try:
                if isinstance(body, str):
                    body_dict = json.loads(body)
                else:
                    body_dict = body
                claude_json = self._convert_copilot_json_to_claude(body_dict)
            except:
                claude_json = {"error": "响应解析失败"}
        
        # 构造Claude响应头
        claude_headers = {
            'Content-Type': 'application/json',
            'Server': 'Claude-Code-Proxy/1.0',
            'Access-Control-Allow-Origin': '*',
            'Access-Control-Allow-Methods': 'GET, POST, OPTIONS',
            'Access-Control-Allow-Headers': 'Content-Type, Authorization'
        }
        
        return web.json_response(
            claude_json,
            status=status,
            headers=claude_headers
        )
    
    def _convert_sse_to_claude_json(self, events: List[Dict[str, Any]]) -> Dict[str, Any]:
        """将SSE事件转换为Claude JSON格式"""
        content_parts = []
        model = "claude-sonnet-4-20250514"
        usage = {"input_tokens": 0, "output_tokens": 0}
        
        for event in events:
            event_type = event.get('type', '')
            data = event.get('data', {})
            
            if event_type == 'content_block_delta':
                delta = data.get('delta', {})
                if delta.get('type') == 'text_delta':
                    content_parts.append(delta.get('text', ''))
            
            elif 'usage' in data:
                usage.update(data['usage'])
            
            elif 'model' in data:
                model = self.transformer._map_model_copilot_to_claude(data['model'])
        
        # 构造Claude响应格式
        claude_response = {
            "id": f"msg_{int(time.time())}",
            "type": "message",
            "role": "assistant", 
            "model": model,
            "content": [
                {
                    "type": "text",
                    "text": "".join(content_parts)
                }
            ],
            "stop_reason": "end_turn",
            "stop_sequence": None,
            "usage": usage
        }
        
        return claude_response
    
    def _convert_copilot_json_to_claude(self, copilot_json: Dict[str, Any]) -> Dict[str, Any]:
        """将Copilot JSON格式转换为Claude格式"""
        # 提取内容
        choices = copilot_json.get('choices', [])
        content = ""
        if choices:
            choice = choices[0]
            message = choice.get('message', {})
            content = message.get('content', '')
        
        # 构造Claude响应
        claude_response = {
            "id": f"msg_{int(time.time())}",
            "type": "message", 
            "role": "assistant",
            "model": self.transformer._map_model_copilot_to_claude(copilot_json.get('model', 'claude-sonnet-4')),
            "content": [
                {
                    "type": "text",
                    "text": content
                }
            ],
            "stop_reason": "end_turn",
            "stop_sequence": None,
            "usage": copilot_json.get('usage', {})
        }
        
        return claude_response
    
    async def _create_error_response(self, error_msg: str) -> Response:
        """创建错误响应"""
        error_response = {
            "error": {
                "type": "api_error",
                "message": error_msg
            }
        }
        
        return web.json_response(
            error_response,
            status=500,
            headers={
                'Content-Type': 'application/json',
                'Access-Control-Allow-Origin': '*'
            }
        )
    
    async def health_check(self, request: Request) -> Response:
        """健康检查端点"""
        health_info = {
            "status": "healthy",
            "timestamp": time.time(),
            "version": "1.0.0"
        }
        return web.json_response(health_info)
    
    async def get_status(self, request: Request) -> Response:
        """获取代理状态"""
        status_info = {
            "proxy_status": "running",
            "listen_address": f"{self.listen_host}:{self.listen_port}",
            "target_url": self.copilot_base_url,
            "transformer_loaded": self.transformer is not None,
            "session_active": self.session is not None and not self.session.closed
        }
        return web.json_response(status_info)
    
    async def cleanup(self):
        """清理资源"""
        if self.session and not self.session.closed:
            await self.session.close()
            self.logger.info("HTTP会话已关闭")


async def main():
    """主函数"""
    proxy = MappingProxy()
    
    try:
        runner = await proxy.start_server()
        
        print("="*60)
        print("🚀 Claude Code <-> Copilot 映射代理服务器已启动")
        print("="*60)
        print(f"📡 监听地址: http://{proxy.listen_host}:{proxy.listen_port}")
        print(f"🔄 转发目标: {proxy.copilot_base_url}")
        print("📋 可用端点:")
        print(f"  • POST /v1/messages - Claude API转换")
        print(f"  • GET /health - 健康检查")
        print(f"  • GET /status - 状态信息")
        print("\n💡 使用方法:")
        print(f"  将Claude Code客户端配置为使用: http://{proxy.listen_host}:{proxy.listen_port}")
        print("\n按 Ctrl+C 停止服务器")
        print("="*60)
        
        # 保持服务器运行
        try:
            while True:
                await asyncio.sleep(1)
        except KeyboardInterrupt:
            print("\n🛑 收到停止信号，正在关闭服务器...")
        
    except Exception as e:
        print(f"❌ 服务器启动失败: {e}")
    
    finally:
        await proxy.cleanup()
        print("✅ 服务器已停止")


if __name__ == "__main__":
    asyncio.run(main())